from ast import mod
from docplex.mp.model import Model
from network import QKDNetwork
from log import addContent

### return 流：[[1,2,3], [4,5,6]] SD对：[[1,3], [4,6]]
def getRoundingX(net: QKDNetwork, x_ret, added_x: list):
    # 此处的n的边容量是1-x_ret的值
    # 使用最短路径算法求解最小权和的路径，并添加到已有的模型中（令x[m][n][i][j]=1）
    length = len(net.nodeList) + 1

    # sd对mn，表示业务量
    sd = []
    for i in range(1, length):
        for j in range(1, length):
            if i != j:
                sd.append((i, j))
                
    total_max_flow_nodes = []
    
    sd_mn = []
    
    for sd_p in sd:
        m, n = sd_p                                 # --> m, n
        contain_flag = False
        for x_tmp in added_x:
            if(x_tmp[0] == m and x_tmp[1] == n):
                contain_flag = True
        if(contain_flag):
            continue
        flowPaths = net.getAllPathFromDFS(m, n, 5)
        single_max_flow = 0
        single_max_flow_nodes = []
        for sdSinglePath in flowPaths:
            # 求解单条流的瓶颈大小
            max_flow = 1
            for i_start in range(len(sdSinglePath) - 1):
                i = sdSinglePath[i_start]
                j = sdSinglePath[i_start + 1]
                if (x_ret[m][n][i][j] < max_flow):
                    max_flow = x_ret[m][n][i][j]
            if(max_flow > single_max_flow):
                single_max_flow = max_flow
                single_max_flow_nodes = sdSinglePath
        flag = False
        if(single_max_flow > 0):
            flag = True
        # if(len(added_x) > (length - 1) * (length - 2) / 2):
        #     flag = True
        if(flag or len(total_max_flow_nodes) == 0):
            total_max_flow_nodes.append(single_max_flow_nodes)
            sd_mn.append([m, n])
                
    return total_max_flow_nodes, sd_mn

def solver(n: QKDNetwork, _, params: dict):
    
    net = n
    length = len(net.nodeList) + 1
    model = Model(name='LP')

    # x_{ij}^{mn}
    # c
    x = {}
    x_ret = {}
    for i in range(length):
        x[i] = {}
        x_ret[i] = {}
        for j in range(length):
            x[i][j] = {}
            x_ret[i][j] = {}
            for k in range(length):
                x[i][j][k] = {}
                x_ret[i][j][k] = {}
                for l in range(length):
                    x[i][j][k][l] = model.continuous_var(
                        lb=0, ub=1, name=f'x_{i}_{j}_{k}_{l}')
                    x_ret[i][j][k][l] = 0

    # sd对mn，表示业务量
    sd = []
    for i in range(1, length):
        for j in range(1, length):
            if i != j:
                sd.append((i, j))

    # p_{ij}^{mnvt}
    # 修复路径所占的比例，mn索引到对应的SD对中，vt索引到SD路由方案上的某条边，ij索引到被分流的边
    p = {}
    p_ret = {}
    for i in range(length):
        p[i] = {}
        p_ret[i] = {}
        for j in range(length):
            p[i][j] = {}
            p_ret[i][j] = {}
            for k in range(length):
                p[i][j][k] = {}
                p_ret[i][j][k] = {}
                for l in range(length):
                    p[i][j][k][l] = {}
                    p_ret[i][j][k][l] = {}
                    for m in range(length):
                        p[i][j][k][l][m] = []
                        p_ret[i][j][k][l][m] = []
                        for q in range(length):
                            p[i][j][k][l][m].append(model.continuous_var(
                                lb=0, ub=1, name=f'p_{i}_{j}_{k}_{l}_{m}_{q}'))
                            p_ret[i][j][k][l][m].append(0)

    # 对偶约束
    lambda_1 = {}
    mu_1 = {}
    for i in range(length):
        lambda_1[i] = {}
        mu_1[i] = {}
        for j in range(length):
            lambda_1[i][j] = []
            mu_1[i][j] = []
            for k in range(length):
                lambda_1[i][j].append(model.continuous_var(
                    lb=0, name=f'lambda_{i}_{j}_{k}'))
                mu_1[i][j].append(model.continuous_var(
                    lb=0, name=f'mu_{i}_{j}_{k}'))
    alpha = model.continuous_var(name='Alpha')

    # 约束1️⃣
    for i in net.nodeList:
        nodeID = i.nodeID                      # --> i
        for adjacentNodeID in i.adjacentNodes:  # --> j
            dualFlowFraction = model.linear_expr()
            dualFlowFraction -= net.getLinkCapacity(
                nodeID, adjacentNodeID) * alpha
            for sd_p in sd:
                sd_start, sd_end = sd_p        # --> m, n
                if sd_start == sd_end:
                    continue
                dualCapacity = model.linear_expr()
                # m --> sd_start
                # n --> sd_end
                # i --> nodeID
                # j --> adjacentNodeID
                dualCapacity += lambda_1[nodeID][adjacentNodeID][sd_start] + mu_1[nodeID][adjacentNodeID][sd_end]
                for tmp_m in net.nodeList:
                    tmp_m_id = tmp_m.nodeID                 # --> u
                    for tmp_n_id in tmp_m.adjacentNodes:    # --> v
                        dualCapacity -= p[sd_start][sd_end][tmp_m_id][tmp_n_id][nodeID][adjacentNodeID]
                for t in net.getNode(sd_start).adjacentNodes:
                    dualFlowFraction += lambda_1[nodeID][adjacentNodeID][sd_start] * \
                        net.getLinkCapacity(sd_start, t)
                for s in net.getNode(sd_end).adjacentNodes:
                    dualFlowFraction += mu_1[nodeID][adjacentNodeID][sd_end] * \
                        net.getLinkCapacity(s, sd_end)
                model.add_constraint(dualCapacity >= 0)
            model.add_constraint(dualFlowFraction <= 0)

    # 约束2️⃣
    # 流约束
    # x_{ij}^{mn} 从业务量起点出发，到终点结束，沿着某一条路径，出点为1，汇点为1
    for sd_p in sd:
        sd_start, sd_end = sd_p
        if sd_start == sd_end:   # --> m,n
            continue
        for i in net.nodeList:
            nodeID = i.nodeID    # --> i
            outFlow = model.linear_expr()
            if nodeID == sd_start:
                outFlow -= 1
            if nodeID == sd_end:
                outFlow += 1
            for adjacentNodeID in i.adjacentNodes:
                if adjacentNodeID == nodeID:
                    continue
                outFlow += x[sd_start][sd_end][nodeID][adjacentNodeID]
                outFlow -= x[sd_start][sd_end][adjacentNodeID][nodeID]
            model.add_constraint(outFlow == 0)

    # 约束3️⃣
    # 修复路径约束
    # 修复路径的流，从边起点到边终点
    for sd_p in sd:
        sd_start, sd_end = sd_p                             # --> m, n
        if sd_start == sd_end:
            continue
        for i in net.nodeList:
            nodeID = i.nodeID                               # --> u
            for adjacentNodeID in i.adjacentNodes:          # --> v
                model.add_constraint(p[sd_start][sd_end][nodeID][adjacentNodeID][nodeID][adjacentNodeID] <= 1 - x[sd_start][sd_end][nodeID][adjacentNodeID])
                for iNode in net.nodeList:
                    nodeI = iNode.nodeID                    # --> i
                    for jNode in iNode.adjacentNodes:       # --> j
                        model.add_constraint(p[sd_start][sd_end][nodeID][adjacentNodeID][nodeI][jNode] <= x[sd_start][sd_end][nodeID][adjacentNodeID])
                    flow = model.linear_expr()
                    for tmp_out in net.getNode(nodeI).adjacentNodes:
                        flow += p[sd_start][sd_end][nodeID][adjacentNodeID][nodeI][tmp_out]
                    for tmp_in in net.getNode(nodeI).adjacentNodes:
                        flow -= p[sd_start][sd_end][nodeID][adjacentNodeID][tmp_in][nodeI]
                    if nodeI == nodeID:
                        flow -= x[sd_start][sd_end][nodeID][adjacentNodeID]
                    if nodeI == adjacentNodeID:
                        flow += x[sd_start][sd_end][nodeID][adjacentNodeID]
                    model.add_constraint(flow == 0)
    
    model.minimize(alpha)

    max_flow_nodes = []
    added_mn = []    # 已增加的SD对
    alpha_ret = 0

    while True:
        # 增加路径流到约束中
        if(len(max_flow_nodes) != 0):
            for item in max_flow_nodes:
                l = len(item)
                start_m = item[0]
                end_n = item[l - 1]
                for tmp_start_node in net.nodeList:
                    tmpNodeID = tmp_start_node.nodeID                       # --> i
                    for tmpAdjacentNodeID in tmp_start_node.adjacentNodes:  # --> j
                        if(is_consecutive(item, tmpNodeID, tmpAdjacentNodeID)):
                            model.add_constraint(x[start_m][end_n][tmpNodeID][tmpAdjacentNodeID] == 1)
                        else:
                            model.add_constraint(x[start_m][end_n][tmpNodeID][tmpAdjacentNodeID] == 0)
        # 模型求解
        sol = model.solve()
        # model.export_as_lp(f"lp/model.lp")
        if sol:
            alpha_ret = sol.get_objective_value()
            for i in range(length):
                for j in range(length):
                    for k in range(length):
                        for l in range(length):
                            x_ret[i][j][k][l] = 0
                            value = sol.get_value(x[i][j][k][l])
                            if value != 0:
                                # print(f"x[{i}][{j}][{k}][{l}]", value)
                                x_ret[i][j][k][l] = value
            for i in range(length):
                for j in range(length):
                    for k in range(length):
                        for l in range(length):
                            for m in range(length):
                                for q in range(length):
                                    p_ret[i][j][k][l][m][q] = 0
                                    value = sol.get_value(p[i][j][k][l][m][q])
                                    if value != 0:
                                        # print(f"p[{i}][{j}][{k}][{l}][{m}][{q}]", value)
                                        p_ret[i][j][k][l][m][q] = value
            max_flow_nodes, sd_mn = getRoundingX(n, x_ret, added_mn)
            # print(f"add: {max_flow_nodes}")
            if (len(max_flow_nodes) == 0 and len(added_mn) == (length - 1) * (length - 2)):
                break
            else:
                # addContent("log.txt", f"add_sd_pair {max_flow_nodes}")
                added_mn.extend(sd_mn)
        else:
            print("求解失败")
            print(model.get_solve_status())
            print(model.get_solve_details())
            break

    # print(x_ret)
    
    return x_ret, p_ret, alpha_ret


def is_consecutive(arr, a, b):
    if a in arr and b in arr:
        idx_a = arr.index(a)
        idx_b = arr.index(b)
        if idx_b - idx_a == 1:
            return True
    return False
